
import os

from PIL import Image
from lxml import etree
from xml.etree import ElementTree
from xml.etree.ElementTree import Element, SubElement


'''
将缩略图的标注修改高清图的标注
'''


ENCODE_METHOD = 'utf-8'

base_dir = 'd:/image_augmentation/'
image_dir = base_dir + 'images/'
annotation_bak_dir = base_dir + 'annotations.bak/'
annotation_dir = base_dir + 'annotations/'


for p_dir,dirs,files in os.walk(image_dir):
    parser = etree.XMLParser(encoding=ENCODE_METHOD)
    for name in files:
        id = os.path.splitext(name)[0]
        xml_path = annotation_bak_dir + id + '.xml'
        if os.path.exists(xml_path):

            img = Image.open(p_dir + name )
            new_width,new_heigth = img.size

            tree = ElementTree.parse(xml_path, parser=parser)
            root = tree.getroot()
            widthElement = root.find('size/width')
            heightElement = root.find('size/height')

            scal_width = new_width / int(widthElement.text)
            scal_height = new_heigth / int(heightElement.text)

            # size
            widthElement.text = str(new_width)
            heightElement.text = str(new_heigth)

            #bndbox
            for object in root.iter('object'):
                xmin = object.find('bndbox/xmin')
                xmin.text = str(round(int(xmin.text) * scal_width))

                xmax = object.find('bndbox/xmax')
                xmax.text = str(round(int(xmax.text)*scal_width))

                ymin = object.find('bndbox/ymin')
                ymin.text = str(round(int(ymin.text) * scal_height))

                ymax = object.find('bndbox/ymax')
                ymax.text = str(round(int(ymax.text) * (scal_height)))

            tree.write(annotation_dir+id+'.xml',encoding=ENCODE_METHOD,xml_declaration=True)

